import ipaddress
import csv
import pandas as pd
import itertools  # 引入itertools用于高效连接迭代器


# ==============================================================================
#  优化的功能函数 (Optimized Function Definitions)
# ==============================================================================

def stream_ips_from_txt(filepath):
    """
    [优化] 从文本文件中以流式（生成器）方式逐行读取IP。
    使用 'yield' 关键字，使其成为一个生成器，不一次性加载所有内容。
    """
    print(f"-> [流式] 正在从 {filepath} 读取IP地址...")
    try:
        with open(filepath, 'r', encoding='utf-8') as file:
            count = 0
            for line in file:
                line = line.strip()
                if line:
                    yield line
                    count += 1
            print(f"   已准备好从文件流式传输 {count} 个条目。")
    except FileNotFoundError:
        print(f"!! 警告: 文件未找到 -> {filepath}。将跳过此源。")
        # 使用空的生成器表达式来返回一个空的迭代器
        return
        yield


def stream_ips_from_excel(filepath, column_name='ip'):
    """
    [优化] 从Excel文件中以流式方式读取IP。
    """
    print(f"-> [流式] 正在从 {filepath} (列: '{column_name}') 读取IP地址...")
    try:
        df = pd.read_excel(filepath)
        if column_name in df.columns:
            count = len(df[column_name].dropna())
            print(f"   已准备好从Excel流式传输 {count} 个额外IP地址。")
            yield from df[column_name].dropna().astype(str)
        else:
            print(f"!! 警告: 文件中未找到名为 '{column_name}' 的列 -> {filepath}")
    except FileNotFoundError:
        print(f"!! 警告: 文件未找到 -> {filepath}。将跳过此源。")
    except Exception as e:
        print(f"!! 读取 {filepath} 文件时出错: {e}")
    return
    yield


def stream_ips_for_asns(asn_list, parquet_path):
    """
    [核心优化] 根据ASN列表，选择性地从Parquet文件中查询并流式传输IP段。
    使用pyarrow的过滤器(predicate pushdown)来避免加载整个文件。
    """
    if not asn_list:
        return
        yield

    print(f"-> [优化查询] 正在从 '{parquet_path}' 为 {len(asn_list)} 个ASN筛选IP...")

    # 使用集合(set)进行过滤，对于大量ASN来说效率更高
    asn_set = set(asn_list)
    # 创建过滤器，这是最重要的性能改进
    filters = [('asn', 'in', asn_set)]

    try:
        # pd.read_parquet 使用 pyarrow 引擎时，可以传递 filters
        # 这将导致只从磁盘读取符合条件的行，极大地减少了内存使用和读取时间
        df_filtered = pd.read_parquet(parquet_path, columns=['network'], filters=filters)

        count = len(df_filtered)
        print(f"   查询完成。仅加载了与ASN匹配的 {count} 行数据。")

        # 报告哪些ASN没有在文件中找到
        found_asns = set(pd.read_parquet(parquet_path, columns=['asn'], filters=filters)['asn'].unique())
        failed_asns = list(asn_set - found_asns)

        # 流式传输找到的IP
        yield from df_filtered['network'].dropna()

        # 将失败的ASN作为最后一个元素yield出去，以便主程序捕获
        # 这是一种在生成器中传递“元数据”或“副作用”结果的技巧
        yield {'failed_asns': failed_asns}

    except FileNotFoundError:
        print(f"!! 严重错误: Parquet文件 '{parquet_path}' 未找到。无法查询ASN。")
        yield {'failed_asns': asn_list}
    except Exception as e:
        print(f"!! 读取或筛选Parquet文件时出错: {e}")
        yield {'failed_asns': asn_list}


def merge_cidrs_from_iterator(cidr_iterator):
    """
    [优化] 从一个迭代器接收CIDR，合并IPv4地址。
    内存效率更高，因为它不需要一个巨大的输入列表。
    """
    print(f"-> 正在从数据流中筛选并合并有效的IPv4地址段...")

    ipv4_networks = []
    # 从迭代器中逐个处理项目
    for ip in cidr_iterator:
        if not ip or not isinstance(ip, str):
            continue
        try:
            network = ipaddress.ip_network(ip, strict=False)
            if network.version == 4:
                ipv4_networks.append(network)
        except ValueError:
            pass  # 忽略格式错误

    if not ipv4_networks:
        print("!! 警告: 未找到任何有效的IPv4地址段可供合并。")
        return []

    print(f"   已从流中收集 {len(ipv4_networks)} 个有效IPv4段，正在进行最终合并...")
    merged = list(ipaddress.collapse_addresses(ipv4_networks))
    print(f"   合并完成。最终生成 {len(merged)} 条IPv4路由。")
    return merged


# 其他辅助函数（read_asn_list_from_csv, generate_routeros_script, generate_failed_asn_report）
# 与前一版本基本相同，此处为简洁省略，可直接复用。
# 我们只需要确保在主逻辑中正确调用它们即可。
def read_asn_list_from_csv(filepath):
    asn_list = []
    print(f"-> 正在从 {filepath} 读取ASN列表...")
    try:
        with open(filepath, mode='r', newline='', encoding='utf-8') as csvfile:
            reader = csv.reader(csvfile)
            next(reader)
            for row in reader:
                if row and row[0]:
                    asn_list.append(row[0].strip())
        print(f"   成功读取 {len(asn_list)} 个ASN。")
    except FileNotFoundError:
        print(f"!! 错误: 文件未找到 -> {filepath}。将返回一个空列表。")
    return asn_list


def generate_routeros_script(networks, gateway, output_path):
    print(f"-> 正在生成RouterOS脚本 -> {output_path}")
    try:
        script_lines = ['/ip route']
        for network in networks:
            script_lines.append(f":do {{add dst-address={network} gateway={gateway} }} on-error={{}}")
        with open(output_path, 'w', encoding='utf-8') as file:
            file.write("\n".join(script_lines))
        print(f"   脚本文件已成功保存。")
    except Exception as e:
        print(f"!! 生成脚本时发生错误: {e}")


def generate_failed_asn_report(failed_asns, output_path):
    if not failed_asns:
        print("-> 所有ASN均已成功查询，无需生成失败报告。")
        return
    print(f"-> 正在将 {len(failed_asns)} 个查询失败的ASN导出到报告 -> {output_path}")
    try:
        df_report = pd.DataFrame(sorted(list(failed_asns)), columns=['ASN'])
        df_report.to_excel(output_path, index=False)
        print(f"   失败报告已成功保存。")
    except Exception as e:
        print(f"!! 导出失败报告时出错: {e}")


# ==============================================================================
#  优化的主程序 (Optimized Main Execution)
# ==============================================================================
def main():
    """
    主函数，以流式和优化的方式协调整个流程。
    """
    print("--- 步骤 1: 加载配置 ---")
    ASN_LIST_PATH = 'asn-all.csv'
    GOOGLE_IP_PATH = 'google_address_route.txt'
    OTHER_IP_EXCEL_PATH = 'other-ip.xlsx'
    IPINFO_DB_PATH = 'ipinfo_lite.parquet'
    GATEWAY_ADDRESS = '172.20.63.253'

    OUTPUT_SCRIPT_FILE = 'generated_routeros_routes.rsc'
    OUTPUT_FAILED_ASN_REPORT = 'failed_asns_report.xlsx'
    print("   配置加载完毕。\n")

    print("--- 步骤 2: 准备输入源 ---")
    asn_to_process = read_asn_list_from_csv(ASN_LIST_PATH)
    if not asn_to_process:
        print("\nASN列表为空或读取失败，程序终止。")
        return

    # 将所有IP源（现在是生成器）链接在一起，而不是加载到列表中
    # itertools.chain 是处理迭代器的最高效方式
    all_cidrs_iterator = itertools.chain(
        stream_ips_from_txt(GOOGLE_IP_PATH),
        stream_ips_from_excel(OTHER_IP_EXCEL_PATH)
    )
    print("   外部IP源已准备为数据流。\n")

    print("--- 步骤 3: 处理ASN数据流 ---")
    # 特殊处理ASN流，因为它还会返回失败的ASN列表
    asn_stream_processor = stream_ips_for_asns(asn_to_process, IPINFO_DB_PATH)

    # 分离IP流和失败ASN报告
    failed_asns = []

    # 创建一个中间生成器来捕获 failed_asns
    def asn_ip_extractor(stream):
        nonlocal failed_asns
        for item in stream:
            if isinstance(item, dict) and 'failed_asns' in item:
                failed_asns = item['failed_asns']
            else:
                yield item

    asn_ip_iterator = asn_ip_extractor(asn_stream_processor)

    # 将ASN的IP流也加入到总的迭代器中
    final_iterator = itertools.chain(all_cidrs_iterator, asn_ip_iterator)
    print("   ASN数据流已连接。\n")

    print("--- 步骤 4: 从数据流合并CIDR ---")
    merged_ipv4_networks = merge_cidrs_from_iterator(final_iterator)
    print("\n")

    print("--- 步骤 5: 生成输出文件 ---")
    if merged_ipv4_networks:
        generate_routeros_script(merged_ipv4_networks, GATEWAY_ADDRESS, OUTPUT_SCRIPT_FILE)
    else:
        print("-> 没有有效的IPv4网络可供生成脚本，已跳过。")

    # generate_failed_asn_report(failed_asns, OUTPUT_FAILED_ASN_REPORT)

    print("\n🎉 全部操作完成！")


if __name__ == "__main__":
    main()